package com.lin.linchatgpt.controller;

import cn.dev33.satoken.annotation.SaCheckLogin;
import cn.dev33.satoken.stp.StpUtil;
import cn.hutool.core.util.ObjUtil;
import com.lin.linchatgpt.entity.Image;
import com.lin.linchatgpt.entity.dto.MidjourneyDTO;
import com.lin.linchatgpt.entity.dto.ResponseMidjourneyDTO;
import com.lin.linchatgpt.midjourney.service.MidjourneyService;
import com.lin.linchatgpt.service.ImageService;
import com.lin.linchatgpt.utils.LinResult;
import com.lin.linchatgpt.ws.GenerateImageMessage;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.BeanUtils;
import org.springframework.web.bind.annotation.*;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Copyright(C),2022年-2022,霖霖
 *
 * @author 霖霖
 * @version 1.0
 * @date 2023/7/7 14:12
 * @Description 图片生成接口
 */
@Slf4j
@RestController
@RequestMapping("lin_chat/generate_image")
public class GenerateImageController {
    @Resource
    private MidjourneyService midjourneyService;

    @Resource
    private ImageService imageService;

    @Resource
    private GenerateImageMessage generateImageMessage;

    private static final Map<String, Image> map = new ConcurrentHashMap<>();

    /**
     * @return com.lin.linchatgpt.utils.LinResult
     * @description 图片生成接口
     * @author 霖霖
     * @date 2023/7/9 9:27
     * @params [java.lang.String, int, int]
     */
    @GetMapping("midjourney")
    @SaCheckLogin
    public LinResult generateImage(@RequestParam("prompt") String prompt,
                                   @RequestParam("modeType") int modelType,
                                   @RequestParam int generateSpeed) {
        MidjourneyDTO midjourneyImage = midjourneyService.getMidjourneyImage(prompt, modelType, generateSpeed);
        if (ObjUtil.isNull(midjourneyImage)) {
            return LinResult.error().message("图片生成失败");
        }
        String loginId = StpUtil.getLoginIdAsString();
        Image image = new Image();
        image.setModel(modelType);
        image.setUserId(loginId);
        image.setPrompt(prompt);


        map.put(midjourneyImage.getTaskId(), image);
        return LinResult.success().data("item", midjourneyImage);
    }


    /**
     * @description 异步接收图片生成结果，将结果存入数据库，同时发送消息给前端，前端接收到消息后，调用接口获取图片信息，展示图片
     * @author 霖霖
     * @date 2023/7/9 9:27
     * @params [com.lin.linchatgpt.entity.dto.ResponseMidjourneyDTO]
     */
    @RequestMapping("receive")
    public void receiveImage(@RequestBody ResponseMidjourneyDTO responseMidjourneyDTO) {
        Image image = map.get(responseMidjourneyDTO.getTaskId());
        log.info("接收的数据为：{}", responseMidjourneyDTO);
        if (ObjUtil.isNull(responseMidjourneyDTO)) {
            responseMidjourneyDTO.setSuccess(false);
            generateImageMessage.sendOneMessage(image.getUserId(), responseMidjourneyDTO);
            log.error("图片接收失败");
        } else {
            if (responseMidjourneyDTO.isSuccess()) {
                log.info("图片接收成功的数据为：{}", responseMidjourneyDTO);
                BeanUtils.copyProperties(responseMidjourneyDTO, image);
                imageService.mySave(image);
                map.remove(responseMidjourneyDTO.getTaskId());
                generateImageMessage.sendOneMessage(image.getUserId(), responseMidjourneyDTO);
            } else {
                log.error("图片接收失败");
            }

        }
    }

    @GetMapping("get_images")
    @SaCheckLogin
    public LinResult getImages() {

        return LinResult.success().data("list", imageService.getImageChat());
    }

    @GetMapping("handler_image")
    @SaCheckLogin
    public LinResult getImages(@RequestParam("action") String action,@RequestParam("imageId") String imageId) {

        MidjourneyDTO midjourneyDTO = midjourneyService.handlerImage(action, imageId);
        if (ObjUtil.isNull(midjourneyDTO)) {
            return LinResult.error().message("图片生成失败");
        }

        String loginId = "264ae6a2a57b4b79aaa0908e7271a807";
        Image image = new Image();
        image.setModel(1);
        image.setUserId(loginId);
        image.setPrompt(action);


        map.put(midjourneyDTO.getTaskId(), image);
        return LinResult.success().message("操作成功").data("item", midjourneyDTO);
    }

}
